import math
import Logger_config
from Logger_config import setup_logger

class TimePredictor:
    """
    模型响应时间预测器
    
    功能：
    1. 根据输入提示词长度和任务类型预测模型响应时间
    2. 支持文档检查、文档优化、文档摘要总结三种任务类型
    
    参数说明：
    - prompt_length: 输入提示词长度，必须为非负整数
    - task_type: 任务类型，可选值：'check', 'optimize', 'summarize'
    
    返回值：
    - 预测的模型响应时间，单位：秒
    
    异常处理：
    - ValueError: 当输入参数无效时抛出异常
    
    计算规则：
    1. 基准长度2500，小于2500按2500计算
    2. 文档检查/摘要：每增加10倍长度，时间增加1.5倍
    3. 文档优化：每增加1倍长度，时间增加1倍
    """

    def __init__(self):
        """初始化日志记录器"""
        self.logger = setup_logger()

    def predict_time(self, prompt_length: int, task_type: str) -> float:
        """
        预测模型响应时间
        
        参数：
        prompt_length (int): 输入提示词长度
        task_type (str): 任务类型，可选值：'check', 'optimize', 'summarize'
        
        返回：
        float: 预测的响应时间（秒）
        
        异常：
        ValueError: 当输入参数无效时抛出
        """
        # 参数校验
        if prompt_length < 0:
            self.logger.error("输入提示词长度不能为负数")
            raise ValueError("输入提示词长度不能为负数")
        
        if task_type not in ['check', 'optimize', 'summarize']:
            self.logger.error(f"无效的任务类型：{task_type}")
            raise ValueError(f"无效的任务类型：{task_type}")

        # 设置基准值
        base_length = 2500
        if prompt_length < base_length:
            effective_length = base_length
            self.logger.debug(f"输入长度小于基准值，使用基准长度{base_length}")
        else:
            effective_length = prompt_length

        # 计算响应时间
        try:
            if task_type == 'check' or task_type == 'summarize':
                # 文档检查和摘要总结逻辑
                if effective_length == base_length:
                    time = 30
                else:
                    # 计算10的幂次数
                    power = math.log10(effective_length / base_length)
                    time = 30 * (1.5 ** power)
                    # 对计算结果向上取整
                    time = math.ceil(time)
            
            elif task_type == 'optimize':
                # 文档优化逻辑
                if effective_length == base_length:
                    time = 60
                else:
                    # 计算2的幂次数（原始值，不取整）
                    power = math.log2(effective_length / base_length)
                    # 直接使用原始power值计算时间
                    time = 60 * (2 ** power)
                    # 对计算结果向上取整
                    time = math.ceil(time)

            return round(time, 0)  # 四舍五入不保留两位小数

        except Exception as e:
            self.logger.error(f"计算响应时间时发生错误：{str(e)}", exc_info=True)
            raise

if __name__ == "__main__":
    # 初始化日志配置
    Logger_config.setup_logger()
    predictor = TimePredictor()
    
    # 测试用例集合
    test_cases = [
        # 基准值测试
        (2500, 'check', 30),
        (2000, 'optimize', 60),  # 小于基准值
        (2500, 'summarize', 30),
        
        # 文档检查/摘要测试
        (25000, 'check', 45),   # 10倍基准值（10^1）
        (250000, 'summarize', 67.5),  # 100倍基准值（10^2）
        
        # 文档优化测试
        (5000, 'optimize', 120),     # 2倍基准值（2^1）
        (10000, 'optimize', 240),    # 4倍基准值（2^2）
        (20000, 'optimize', 480),    # 8倍基准值（2^3）
        
        # 边界值测试
        (0, 'check', 30),          # 零值测试
        (1000000, 'optimize', 3840)  # 超大值测试（2^8）
    ]
    
    # 执行测试用例
    for length, task_type, expected in test_cases:
        try:
            result = predictor.predict_time(length, task_type)
            print(f"输入: {length}, 任务: {task_type}, 实际: {result}, 预期: {expected}, {'✓' if result == expected else '✗'}")
        except Exception as e:
            print(f"输入: {length}, 任务: {task_type}, 执行失败: {str(e)}")
    
    # 非法输入测试
    invalid_cases = [
        (-100, 'check'),     # 负数长度
        (2500, 'invalid')     # 无效任务类型
    ]
    
    for length, task_type in invalid_cases:
        try:
            result = predictor.predict_time(length, task_type)
            print(f"输入: {length}, 任务: {task_type}, 不应到达此行")
        except ValueError as e:
            print(f"输入: {length}, 任务: {task_type}, 拒绝: {str(e)} ✓")
